/*
	Copyright (C) 2010 Salil Bhagurkar

	This file is part of k3

	k3 is free software: you can redistribute it and/or modify
	it under the terms of the GNU Lesser General Public License as published by
	the Free Software Foundation, either version 3 of the License, or
	(at your option) any later version.

	k3 is distributed in the hope that it will be useful,
	but WITHOUT ANY WARRANTY; without even the implied warranty of
	MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
	GNU Lesser General Public License for more details.

	You should have received a copy of the GNU Lesser General Public License
	along with k3.  If not, see <http://www.gnu.org/licenses/>.
*/

#include <shared.h>

/**
TODO:
Sanity checks
*/

//#define DEBUG_MALLOC


#define DESC_MAGIC ('M' | 'H' << 8 | 'D' << 16 | 'R' << 24)
#define DESC_STATUS_ALLOC 1
#define DESC_STATUS_RESERVED 2
#define DESC_STATUS_FREE 0

#define ALLOC_UNIT 4

#define ZONE_ID_LEN 8

/**
Descriptor for memory
*/
struct desc {
	u32 magic_1;
	char zone_id[ZONE_ID_LEN];
	u8 status;
	struct desc *next, *prev;
	void *start;
	void *end;
	u32 magic_2;
};


struct malloc_system {
	void *start;
	void *end;
	struct desc *head;
};



#define desc_size(desc) ((u32)(desc->end) - (u32)(desc->start))

/**
 * Structure to store the currently initialized malloc_get info
 */
static struct malloc_system ms_kern;


/**
 * This adds a zone that is allocated to the list of zones
 * It should be called only after the malloc_system is valid
 */
void malloc_add_reserved_zone(void *start, void *end, char *zone_id)
{
	struct desc *head;
	head = ms_kern.head;
	struct desc *new_desc = (struct desc *)malloc_get(sizeof(struct desc), "mdesc");
	new_desc->start = start;
	new_desc->end = end;
	new_desc->magic_1 = DESC_MAGIC;
	new_desc->magic_2 = DESC_MAGIC;
	new_desc->next = head->next;
	new_desc->prev = head;
	new_desc->status = DESC_STATUS_RESERVED;
	head->next = new_desc;
	memset(new_desc->zone_id, 0, ZONE_ID_LEN);
	string_copy(new_desc->zone_id, zone_id, ZONE_ID_LEN);
}

void malloc_add_free_zone(void *start, void *end, char *zone_id)
{
	struct desc *head;
	head = ms_kern.head;
	struct desc *new_desc = (struct desc *)malloc_get(sizeof(struct desc), "mdesc");
	new_desc->start = start;
	new_desc->end = end;
	new_desc->magic_1 = DESC_MAGIC;
	new_desc->magic_2 = DESC_MAGIC;
	new_desc->next = head->next;
	new_desc->prev = head;
	new_desc->status = DESC_STATUS_FREE;
	head->next = new_desc;
	memset(new_desc->zone_id, 0, ZONE_ID_LEN);
	string_copy(new_desc->zone_id, zone_id, ZONE_ID_LEN);
}


/**
Initiate a header in a zone marking it as malloc_free/allocated
*/
static struct desc *malloc_init_zone(struct desc *desc, void *start, void *end, u8 status, struct desc *nextzone, struct desc *prevzone, char *zone_id)
{
	//Initialize all parameters
	desc->start = start;
	desc->end = end;
	desc->magic_1 = DESC_MAGIC;
	desc->magic_2 = DESC_MAGIC;
	desc->next = nextzone;
	desc->prev = prevzone;
	desc->status = status;
	memset(desc->zone_id, 0, ZONE_ID_LEN);
	string_copy(desc->zone_id, zone_id, ZONE_ID_LEN);
	return desc;
}


/**
Split a malloc_free zone into a first zone of size 'first_sz' and the remaining zone
*/
static void malloc_split_free_zone(struct desc *zone, u32 first_sz)
{
	dprint(D_FREQ, "zone", "%x with %u", zone, first_sz);
	struct desc *next_zone;
	u32 total_sz;
	total_sz = desc_size(zone);
	//Do not split this zone if the second zone will
	//be useless. In that case, the first zone itself should be allocated
	if((total_sz - first_sz) < (ALLOC_UNIT + sizeof(struct desc)))
		return;
	next_zone = zone->start + first_sz;
	malloc_init_zone(next_zone, next_zone + 1, zone->end, DESC_STATUS_FREE, zone->next, zone, ".right");
	//End of this zone is just before the descriptor of the next zone
	malloc_init_zone(zone, zone + 1, zone->start + first_sz, DESC_STATUS_FREE, next_zone, zone->prev, ".left");
}


/**
Find a zone of size 'sz' in the list of memory descriptors
*/
static struct desc *malloc_find_free_zone(u32 sz)
{
	struct desc *t;
	u32 pad;
	struct malloc_system *ms = &ms_kern;
	for(t = ms->head; t; t = t->next) {
		if(desc_size(t) == sz && t->status == DESC_STATUS_FREE) {
			return t;
		}
	}

	//Find a greater zone and split it
	for(t = ms->head; t; t = t->next) {
		if(desc_size(t) > sz && t->status == DESC_STATUS_FREE) {
			malloc_split_free_zone(t, sz);
			return t;
		}
	}
	return NULL;
}

static void malloc_merge_zone(struct desc *t1, struct desc *t2)
{
	dprint(D_FREQ, "zone", "%x, %x", t1, t2);
	t2->magic_1 = 0; //Invalidate the second zone's header
	t2->magic_2 = 0;
	t1->end = t2->end;
	t1->next = t2->next;
	memset(t1->zone_id, 0, ZONE_ID_LEN);
	string_copy(t1->zone_id, ".merge", ZONE_ID_LEN);
}

#define MAX_COUNT 100

static int malloc_merge_free_zones()
{
	struct desc *t;
	u32 merges = 0;
	struct malloc_system *ms = &ms_kern;
	for(t = ms->head; t;) {
		if(t->status == DESC_STATUS_FREE)
		{
			if(t->next)
			{
				if(t->next->status == DESC_STATUS_FREE)
				{
					malloc_merge_zone(t, t->next);
					merges++;
					//Without moving to the next node, try to merge with the next
					continue;
				}
			}
		}
		t = t->next;
	}
	return merges;
}


void *malloc_get(u32 actual_sz, char *zone_id)
{
	struct desc *zone;
	u32 sz;
	struct malloc_system *ms = &ms_kern;
	sz = actual_sz / ALLOC_UNIT * ALLOC_UNIT + ALLOC_UNIT;
	zone = malloc_find_free_zone(sz);
	if(zone) {
		zone->status = DESC_STATUS_ALLOC;
		memset(zone->zone_id, 0, ZONE_ID_LEN);
		string_copy(zone->zone_id, zone_id, ZONE_ID_LEN);
		dprint(D_FREQ, zone->zone_id, "%x", zone->start);
		return zone->start;
	} else {
		printk("malloc failed! (actual_sz = %u, zone_id = %s)\n", actual_sz, zone_id);
		malloc_verify();
		__malloc_print_info(false, 0);
		while(1);
	}
}

int malloc_free(void *mem)
{
	struct desc *zone;
	struct malloc_system *ms = &ms_kern;
	u32 aligned_addr = (u32)mem;
	zone = (struct desc *)aligned_addr;
	zone--;
	if((zone->magic_1) != DESC_MAGIC || (zone->magic_2) != DESC_MAGIC) {
		dprint(D_ALERT, "zone", "Zone magic is invalid (%x)", mem);
		malloc_verify();
		return ESTATEINV;
	}
	zone->status = DESC_STATUS_FREE;
	while(malloc_merge_free_zones());
	dprint(D_FREQ, zone->zone_id, "%x");
	return 0;
}

/**
 * Reallocate more size for the memory allocation
 */
void *malloc_resize(void *mem, u32 sz)
{
	dprint(D_FREQ, "realloc", "%x to %u", mem, sz);
	void *new_mem;
	struct desc *zone;
	zone = (struct desc *)mem;
	zone--;
	if((zone->magic_1) != DESC_MAGIC || (zone->magic_2) != DESC_MAGIC) {
		dprint(D_ALERT, "zone", "Zone magic is invalid (%x)", mem);
		malloc_verify();
	}
	if(desc_size(zone) > 2000) {
		__malloc_print_info(true, 0);
	}
	new_mem = malloc_get(sz, zone->zone_id);
	memcpy(new_mem, mem, desc_size(zone));
	zone->status = DESC_STATUS_FREE;
	while(malloc_merge_free_zones());
	return new_mem;
}

static void print_memory(u32 sz)
{
	if(sz > 1048576) {
		printk("%u MiB", sz / 1048576);
	} else if(sz > 1024) {
		printk("%u KiB", sz / 1024);
	} else {
		printk("%u B", sz);
	}
}

void __malloc_print_info(bool print_list, u32 limit_records)
{
	struct desc *t;
	u32 sz;
	int i = 0;
	struct malloc_system *ms = &ms_kern;
	u32 used_memory = 0, free_memory = 0, reserved_memory = 0, used_slots = 0, free_slots = 0, reserved_slots = 0;
	for(t = ms->head; t; t = t->next) {
		if(print_list && (limit_records > 0 ? (i++ < limit_records) : true)) {
			printk("%8x - %8x  (%s) is ", t->start, t->end, ((t->magic_1 == DESC_MAGIC) && (t->magic_2 == DESC_MAGIC)) ? "verified" : "invalid ");
			switch(t->status) {
			case DESC_STATUS_ALLOC:
				printk("allocated");
				break;
			case DESC_STATUS_RESERVED:
				printk("reserved ");
				break;
			case DESC_STATUS_FREE:
				printk("free     ");
				break;
			}
			printk(" size = %8x, id = %s\n", desc_size(t), t->zone_id);
		}
		sz = desc_size(t);
		if(t->status == DESC_STATUS_ALLOC) {
			used_memory += sz;
			used_slots++;
		} else if(t->status == DESC_STATUS_RESERVED) {
			reserved_memory += sz;
			reserved_slots++;
		} else if(t->status == DESC_STATUS_FREE) {
			free_memory += sz;
			free_slots++;
		}
	}
	if(malloc_verify() == 0) {
		printk("All signatures verified\n");
	}
	printk("Descriptor size = %x\n", sizeof(struct desc));
	printk("Alloc unit = %u\n", ALLOC_UNIT);
	printk("Used slots: %u (", used_slots);
	print_memory(used_memory);
	printk(")");
	printk("\nReserved slots: %u (", reserved_slots);
	print_memory(reserved_memory);
	printk(")");
	printk("\nTotal used memory: ");
	print_memory(used_memory + reserved_memory);
	printk("\nFree slots: %u (", free_slots);
	print_memory(free_memory);
	printk(")\n");
}

/**
 * Function to verify all magics and all 'end' members
 */
int malloc_verify()
{
	struct desc *t;
	u32 sz;
	int i = 0;
	struct malloc_system *ms = &ms_kern;
	for(t = ms->head; t; t = t->next) {
		sz = desc_size(t);
		if(t->magic_1 != DESC_MAGIC) {
			dprint(D_ALERT, t->zone_id, "Invalid zone found %x", t);
			dprint(D_ALERT, t->prev->zone_id, "Possible overwrite by previous zone %x", t->prev);
			return ESTATEINV;
		} else if(t->magic_2 != DESC_MAGIC) {
			dprint(D_ALERT, t->zone_id, "Invalid zone found %x", t);
			dprint(D_ALERT, t->next->zone_id, "Possible overwrite by next zone %x", t->next);
		}
		if(((u32)t->end) < ((u32)t->start)) {
			dprint(D_ALERT, t->zone_id, "Zone end is before zone start %x", t);
			return ESTATEINV;
		}
	}
	return 0;
}

void malloc_init()
{
	dprint(D_PROG, "malloc", "Initialize");
	ms_kern.start = (void *)arch_kernel_end;
	ms_kern.end = (void *)0x900000;
	ms_kern.head = malloc_init_zone(ms_kern.start, ms_kern.start + sizeof(struct desc), ms_kern.end, DESC_STATUS_FREE, NULL, NULL, ".base");
	arch_register_memory();

	//~ __malloc_print_info(true, 0);
//~
	//~ void *alloc = malloc_get(1, "test");
	//~ alloc = malloc_resize(alloc, 2);
	//~ alloc = malloc_resize(alloc, 3);
	//~ alloc = malloc_resize(alloc, 4);
	//~ alloc = malloc_resize(alloc, 5);
	//~ alloc = malloc_resize(alloc, 6);
//~
	//~ __malloc_print_info(true, 0);

	//~ while(1)
		//~ arch_halt();
}
